# knowledge/processors/kb_document_processor.py

import os
import logging
from datetime import datetime
from typing import List, Dict, Tuple
from llama_index.core import Document

class KBDocumentProcessor:
    """知识库文档处理器"""
    
    def __init__(self):
        self.logger = logging.getLogger("KnowledgeBase.document_processor")
    
    def scan_knowledge_base(self, root_path: str) -> Dict[str, List[str]]:
        """
        扫描知识库目录结构
        
        Args:
            root_path: 知识库根目录
            
        Returns:
            dict: 按文件夹分类的文件路径
        """
        knowledge_structure = {}
        
        if not os.path.exists(root_path):
            self.logger.warning(f"知识库根目录不存在: {root_path}")
            return knowledge_structure
        
        try:
            for folder_name in os.listdir(root_path):
                folder_path = os.path.join(root_path, folder_name)
                if os.path.isdir(folder_path):
                    folder_files = self._scan_folder(folder_path)
                    knowledge_structure[folder_name] = folder_files
                    self.logger.info(f"扫描文件夹 '{folder_name}': 找到 {len(folder_files)} 个文件")
                else:
                    # 如果根目录下有文件，放到 "root" 分类中
                    if self._is_supported_file(folder_path):
                        if "root" not in knowledge_structure:
                            knowledge_structure["root"] = []
                        knowledge_structure["root"].append(folder_path)
            
        except Exception as e:
            self.logger.error(f"扫描知识库目录失败: {e}")
        
        return knowledge_structure
    
    def _scan_folder(self, folder_path: str) -> List[str]:
        """扫描文件夹中的所有支持文件"""
        supported_files = []
        
        try:
            for root, dirs, files in os.walk(folder_path):
                for file in files:
                    file_path = os.path.join(root, file)
                    if self._is_supported_file(file_path):
                        supported_files.append(file_path)
        except Exception as e:
            self.logger.error(f"扫描文件夹失败 {folder_path}: {e}")
        
        return supported_files
    
    def load_documents_with_metadata(self, file_paths: List[str], root_path: str) -> List[Document]:
        """
        加载文档并添加丰富的元数据
        
        Args:
            file_paths: 文件路径列表
            root_path: 知识库根目录
            
        Returns:
            list: 文档对象列表
        """
        documents = []
        
        for file_path in file_paths:
            try:
                # 计算相对路径和文件夹层级
                relative_path = os.path.relpath(file_path, root_path)
                folder_hierarchy = relative_path.split(os.sep)
                
                # 加载文档内容
                content = self._read_file_content(file_path)
                if not content or not content.strip():
                    self.logger.warning(f"⚠️ 跳过空文件: {file_path}")
                    continue
                
                # 创建文档对象
                doc = Document(
                    text=content,
                    metadata={
                        "file_path": file_path,
                        "relative_path": relative_path,
                        "folder_hierarchy": "/".join(folder_hierarchy[:-1]),  # 排除文件名
                        "folder_name": folder_hierarchy[0] if len(folder_hierarchy) > 1 else "root",
                        "file_name": os.path.basename(file_path),
                        "file_extension": os.path.splitext(file_path)[1].lower(),
                        "depth": len(folder_hierarchy) - 1,
                        "last_modified": datetime.fromtimestamp(os.path.getmtime(file_path)).isoformat()
                    }
                )
                documents.append(doc)
                self.logger.info(f"✅ 加载文档: {relative_path} (长度: {len(content)} 字符)")
                
            except Exception as e:
                self.logger.error(f"❌ 加载文档失败 {file_path}: {e}")
        
        self.logger.info(f"文档加载完成: {len(documents)} 个有效文档")
        return documents
    
    def _is_supported_file(self, file_path: str) -> bool:
        """检查是否支持的文件类型"""
        supported_extensions = ['.pdf', '.txt', '.md', '.markdown', '.jsonl', '.doc', '.docx']
        file_ext = os.path.splitext(file_path)[1].lower()
        return file_ext in supported_extensions and os.path.isfile(file_path)
    
    def _read_file_content(self, file_path: str) -> str:
        """读取文件内容 - 修复PDF处理"""
        file_ext = os.path.splitext(file_path)[1].lower()
        
        try:
            if file_ext == '.pdf':
                return self._read_pdf_file_fixed(file_path)
            elif file_ext in ['.txt', '.md', '.markdown']:
                return self._read_text_file(file_path)
            elif file_ext == '.jsonl':
                return self._read_jsonl_file(file_path)
            else:
                # 对于其他格式，尝试文本读取
                try:
                    return self._read_text_file(file_path)
                except:
                    self.logger.warning(f"不支持的文件格式，跳过: {file_path}")
                    return ""
        except Exception as e:
            self.logger.error(f"读取文件失败 {file_path}: {e}")
            return ""
        
    def _read_pdf_file_fixed(self, file_path: str) -> str:
        """读取PDF文件 - 修复版本"""
        try:
            # 首先尝试使用简单的PDF读取器
            from llama_index.core import SimpleDirectoryReader
            pdf_docs = SimpleDirectoryReader(input_files=[file_path]).load_data()
            content = "\n".join([doc.text for doc in pdf_docs if doc.text.strip()])
            
            if content.strip():
                self.logger.info(f"✅ PDF文件读取成功: {file_path} -> {len(content)} 字符")
                return content
            else:
                self.logger.warning(f"⚠️ PDF文件内容为空: {file_path}")
                return ""
                
        except Exception as e:
            self.logger.error(f"❌ PDF文件读取失败 {file_path}: {e}")
            # 尝试回退到基本的PDF读取
            try:
                return self._read_pdf_basic(file_path)
            except:
                return ""
            
    def _read_pdf_basic(self, file_path: str) -> str:
        """基本的PDF读取方法"""
        try:
            import PyPDF2
            content = ""
            with open(file_path, 'rb') as file:
                pdf_reader = PyPDF2.PdfReader(file)
                for page in pdf_reader.pages:
                    text = page.extract_text()
                    if text:
                        content += text + "\n"
            return content
        except Exception as e:
            self.logger.error(f"❌ 基本PDF读取失败 {file_path}: {e}")
            return ""
    
    def _read_text_file(self, file_path: str) -> str:
        """读取文本文件"""
        try:
            with open(file_path, 'r', encoding='utf-8') as f:
                content = f.read()
                if content.strip():
                    return content
                else:
                    self.logger.warning(f"⚠️ 文本文件内容为空: {file_path}")
                    return ""
        except UnicodeDecodeError:
            # 尝试其他编码
            try:
                with open(file_path, 'r', encoding='gbk') as f:
                    content = f.read()
                    if content.strip():
                        return content
                    else:
                        return ""
            except:
                self.logger.error(f"❌ 文本文件编码读取失败: {file_path}")
                return ""
        except Exception as e:
            self.logger.error(f"❌ 文本文件读取失败 {file_path}: {e}")
            return ""
    
    def _read_pdf_file(self, file_path: str) -> str:
        """读取PDF文件"""
        try:
            from knowledge.util.doc_util import load_pdf_documents
            pdf_docs = load_pdf_documents(file_path)
            return "\n".join([doc.text for doc in pdf_docs])
        except ImportError:
            # 回退到简单PDF加载
            from llama_index.core import SimpleDirectoryReader
            pdf_docs = SimpleDirectoryReader(input_files=[file_path]).load_data()
            return "\n".join([doc.text for doc in pdf_docs])
        except Exception as e:
            self.logger.error(f"PDF文件读取失败 {file_path}: {e}")
            return ""
    
    def _read_jsonl_file(self, file_path: str) -> str:
        """读取JSONL文件"""
        try:
            import json
            content_lines = []
            with open(file_path, 'r', encoding='utf-8') as f:
                for line in f:
                    line = line.strip()
                    if line:
                        try:
                            data = json.loads(line)
                            # 提取文本内容
                            if isinstance(data, dict):
                                text = data.get('text', '') or data.get('content', '')
                                if text:
                                    content_lines.append(str(text))
                        except json.JSONDecodeError:
                            continue
            
            content = "\n".join(content_lines)
            if content.strip():
                return content
            else:
                self.logger.warning(f"⚠️ JSONL文件内容为空: {file_path}")
                return ""
        except Exception as e:
            self.logger.error(f"❌ JSONL文件读取失败 {file_path}: {e}")
            return ""